package org.example.ratelimit;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 令牌桶算法
 * token以恒定速率生成
 * lastToken 剩余令牌数量
 * lastTime 上次令牌生成时间
 */
@Slf4j
public class TokenBucketLimiter {

    /**
     * 桶的容量：10
     */
    private static final long capacity = 10;

    /**
     * tokens生成速率
     */
    private static final long rate = 1000;

    /**
     * 剩余tokens
     */
    private static long tokens = 0;

    /**
     * 上次生成的时间
     */
    private static long lastOutTime = System.currentTimeMillis();
    /**
     * true：成功获取令牌，放行
     * false：获取令牌失败，丢弃
     * @param taskId
     * @param applyCount
     * @return
     */
    public synchronized static boolean tryAccquire(long taskId, int applyCount) {
        long nowTime = System.currentTimeMillis();
        long pastTime = nowTime - lastOutTime;

        // 还有上次的容量
        tokens = Math.min(capacity, (tokens + pastTime * rate) / 1000);
        // 令牌是否足够
        if (tokens >= applyCount) {
            tokens = tokens - applyCount;
            lastOutTime = nowTime;
            log.info("新时间区到了，taskId：{}, turns：{}", taskId, applyCount);
            return false;
        } else {
            return true;
        }
    }

    public static void main(String[] args) throws Exception{
        ExecutorService service = Executors.newFixedThreadPool(2);
        CountDownLatch latch = new CountDownLatch(2);

        // 每个线程执行20次
        final int turns = 20;
        final AtomicInteger failed = new AtomicInteger();
        long start = System.currentTimeMillis();
        for (int i = 0; i < 2; i++) {
            service.submit(() -> {
                try {
                    for (int j = 0; j < turns; j++) {
                        long taskId = Thread.currentThread().getId();
                        // 每次获取1个令牌
                        boolean result = tryAccquire(taskId, 1);
                        if (result) {
                            failed.incrementAndGet();
                        }
                        Thread.sleep(200);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    latch.countDown();
                }
            });
        }

        latch.await();
        log.info("总的请求数：{}", turns * 2);
        log.info("被限制的请求数：{}", failed.get());
        log.info("成功的请求数：{}", turns * 2 - failed.get());
        log.info("成功比例：{}", (float)((turns * 2 - failed.get())) / (float) ((turns * 2)));
        log.info("请求总耗时：{}秒", (System.currentTimeMillis() - start) / 1000F);
    }
}
